# 跟踪令牌使用情况

# from langchain.utilities import DuckDuckGoSearchAPIWrapper
from langchain_community.utilities import DuckDuckGoSearchAPIWrapper
from langchain_openai import OpenAI, ChatOpenAI
# from langchain.callbacks import get_openai_callback
from langchain_community.callbacks.manager import get_openai_callback

from openaiConfigurations import openai_api_key, openai_api_base
from langchain.agents import AgentType, initialize_agent, Tool
# from langchain.agents import load_tools
from langchain_community.agent_toolkits.load_tools import load_tools
from langchain.tools import tool
import requests

# llm = OpenAI(openai_api_key = openai_api_key,
#                  openai_api_base = openai_api_base,
#                  n = 2, # 一共生成两个文本
#                  best_of = 2, # 挑选两个最好的返回
#                  model = "text-davinci-002")

# with get_openai_callback() as cb:
#     result = llm("Tell me a joke")
#     print(cb)

"""输出信息如下：
Tokens Used: 42
    Prompt Tokens: 4
    Completions Tokens: 38
Successfully Requests: 1
Total Cost (USD): $0.00084"""

# 顺序跟踪多个调用
# with get_openai_callback() as cb:
#     result = llm("Tell me a joke")
#     result2 = llm("Tell me a joke")
#     print(cb.total_tokens)
# cb.total_cost

llm2 = ChatOpenAI(openai_api_key = openai_api_key,
                 openai_api_base = openai_api_base,
                 temperature = 0)

# DuckDuckGo目前不能在国内访问,下面是DuckDuckGo的配置方法
# search = DuckDuckGoSearchAPIWrapper(timeout=15)
# tools = [
#     Tool(
#         name = "DuckDuckGo Search",
#         func = search.run,
#         description = "Useful for finding information on the internet. Input should be a search query."
#     )
# ]
# agent = initialize_agent(tools, llm2, agent="zero-shot-react-description", verbose=True)

@tool
def bocha_websearch_tool(query: str, count: int = 5) -> str:
    """
    使用Bocha Web Search API 进行网页搜索。

    参数:
    - query: 搜索关键词
    - freshness: 搜索的时间范围
    - summary: 是否显示文本摘要
    - count: 返回的搜索结果数量

    返回:
    - 搜索结果的详细信息，包括网页标题、网页URL、网页摘要、网站名称、网站Icon、网页发布时间等。
    """

    url = 'https://api.bochaai.com/v1/web-search'
    headers = {
        'Authorization': f'Bearer sk-5fb1b126b1ec40c3b654baee5e53235e',  # 请替换为你的API密钥
        'Content-Type': 'application/json'
    }
    data = {
        "query": query,
        "freshness": "noLimit", # 搜索的时间范围，例如 "oneDay", "oneWeek", "oneMonth", "oneYear", "noLimit"
        "summary": True, # 是否返回长文本摘要
        "count": count
    }

    response = requests.post(url, headers=headers, json=data)

    if response.status_code == 200:
        json_response = response.json()
        try:
            if json_response["code"] != 200 or not json_response["data"]:
                return f"搜索API请求失败，原因是: {response.msg or '未知错误'}"
            
            webpages = json_response["data"]["webPages"]["value"]
            if not webpages:
                return "未找到相关结果。"
            formatted_results = ""
            for idx, page in enumerate(webpages, start=1):
                formatted_results += (
                    f"引用: {idx}\n"
                    f"标题: {page['name']}\n"
                    f"URL: {page['url']}\n"
                    f"摘要: {page['summary']}\n"
                    f"网站名称: {page['siteName']}\n"
                    f"网站图标: {page['siteIcon']}\n"
                    f"发布时间: {page['dateLastCrawled']}\n\n"
                )
            return formatted_results.strip()
        except Exception as e:
            return f"搜索API请求失败，原因是：搜索结果解析失败 {str(e)}"
    else:
        return f"搜索API请求失败，状态码: {response.status_code}, 错误信息: {response.text}"

# 创建LangChain工具
bocha_tool = Tool(
    name="BochaWebSearch",
    func=bocha_websearch_tool,
    description="使用Bocha Web Search API 进行搜索互联网网页，输入应为搜索查询字符串，输出将返回搜索结果的详细信息，包括网页标题、网页URL、网页摘要、网站名称、网站Icon、网页发布时间等。"
)

# tools = load_tools(["serpapi", "llm-math"], llm = llm2, serpapi_api_key= serpapi_api_key)
agent = initialize_agent(
    tools=[bocha_tool], 
    llm = llm2, 
    agent = AgentType.ZERO_SHOT_REACT_DESCRIPTION,
    verbose = True
)

# 持续跟踪
with get_openai_callback() as cb:
    response = agent.run(
        "who is Olivia wilde's boyfriend? What is his current age raised to the 0.23 power"
    )
    print(f"total tokens: ", cb.total_tokens)
    print(f"prompt tokens: ", cb.prompt_tokens)
    print(f"completion tokens: ", cb.completion_tokens)
    print(f"total costs: ", cb.total_cost)


'''注释：这段代码无法运行
       主要原因是智能体调用bocha进行搜索时会超出代理的token限制，可以充钱解决。
'''
